pub mod linked_list;

use std::hash::Hash;

use hashbrown::HashMap;
use linked_list::{Index, LinkedList, LinkedListIter};

struct LruNode<T> {
    data: T,
    list_index: Index,
    order_index: usize,
}

pub struct Lru<K: Default + Clone + Eq + Hash, T> {
    map: HashMap<K, Box<LruNode<T>>>,
    orders: Vec<LinkedList<K>>,
    order_capacity: usize,
}

#[allow(dead_code)]
impl<K: Default + Clone + PartialEq + Eq + Hash, T> Lru<K, T> {
    #[inline]
    pub fn with_capacity(nlist: usize, capacity: usize) -> Self {
        assert!(nlist > 0 && capacity > nlist);

        let order_capacity = capacity / nlist;
        let mut orders = Vec::new();
        for _ in 0..nlist {
            orders.push(LinkedList::with_capacity(order_capacity + 2));
        }

        Lru {
            map: HashMap::with_capacity(capacity),
            orders,
            order_capacity,
        }
    }

    // 使用key从 `Lru`中查询, 不改变Lru中的顺序.
    #[inline]
    pub fn peek(&self, key: K) -> Option<&T> {
        self.map.get(&key).map(|n| &n.data)
    }

    #[inline]
    pub fn access(&mut self, key: K) -> Option<&T> {
        let mut node = match self.map.remove(&key) {
            Some(node) => node,
            None => {
                return None;
            }
        };

        let cur_order_index = node.order_index;
        self.orders[cur_order_index].remove(node.list_index);

        // Already at the highest level.
        if cur_order_index == self.orders.len() - 1 {
            node.list_index =
                self.orders[node.order_index].push_head(key.clone());
            self.map.insert(key.clone(), node);
            return self.peek(key.clone());
        }

        // Move to upper order list.
        node.order_index += 1;

        // The upper order list not full.
        if self.orders[cur_order_index + 1].len() < self.order_capacity {
            node.list_index =
                self.orders[node.order_index].push_head(key.clone());
            self.map.insert(key.clone(), node);
            return self.peek(key.clone());
        }

        // The upper level is full, swap.
        let last_key = self.orders[cur_order_index + 1].pop_tail().unwrap();
        let last_node = self.map.get_mut(&last_key).unwrap();
        last_node.order_index -= 1;
        last_node.list_index =
            self.orders[last_node.order_index].push_head(last_key);

        node.list_index = self.orders[node.order_index].push_head(key.clone());
        self.map.insert(key.clone(), node);

        self.peek(key)
    }

    #[inline]
    pub fn remove(&mut self, key: K) -> Option<T> {
        self.map.remove(&key).map(|node| {
            let list_key =
                self.orders[node.order_index].remove(node.list_index);
            assert!(key == list_key);
            node.data
        })
    }

    // 淘汰第一层级order list中的最后一个元素
    #[inline]
    pub fn evict(&mut self) -> Option<T> {
        self.orders[0].pop_tail().map(|key| {
            let node = self.map.remove(&key).unwrap();
            node.data
        })
    }

    // 在第一层的order中添加元素, 如果LRU中有相同key的data, 删除并返回它.
    #[inline]
    pub fn insert(&mut self, key: K, data: T) -> Option<T> {
        let mut old = self.remove(key.clone());
        if old.is_none() && self.orders[0].len() == self.order_capacity {
            old = self.evict();
        }

        let list_index = self.orders[0].push_head(key.clone());
        let node = Box::new(LruNode {
            data,
            list_index,
            order_index: 0,
        });
        self.map.insert(key, node);

        old
    }

    #[inline]
    pub fn len(&self) -> usize {
        let mut res = 0;
        for l in &self.orders {
            res += l.len();
        }

        res
    }

    #[inline]
    pub fn iter(&self) -> LruIter<'_, '_, K, T> {
        let mut base_iters = Vec::with_capacity(self.orders.len());

        for v in self.orders.iter().rev() {
            base_iters.push(v.iter());
        }

        LruIter {
            lru: self,
            base_iters,
            cur_index: 0,
        }
    }
}

pub struct LruIter<'a, 'b, K: Default + Clone + Eq + Hash, T> {
    lru: &'a Lru<K, T>,
    base_iters: Vec<LinkedListIter<'b, K>>,
    cur_index: usize,
}

impl<'a, 'b, K: Default + Clone + Eq + Hash, T> Iterator
    for LruIter<'a, 'b, K, T>
{
    type Item = (K, &'a T);
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let base_iter = &mut self.base_iters[self.cur_index];
            if let Some(key) = base_iter.next() {
                return Some((
                    key.clone(),
                    self.lru.peek(key.clone()).unwrap(),
                ));
            }

            if self.cur_index == self.base_iters.len() - 1 {
                return None;
            }

            self.cur_index += 1;
        }
    }
}

impl<'a, 'b, K: Default + Clone + Eq + Hash, T> DoubleEndedIterator
    for LruIter<'a, 'b, K, T>
{
    fn next_back(&mut self) -> Option<Self::Item> {
        let n_iters = self.base_iters.len();
        if n_iters == 0 {
            return None;
        }

        loop {
            let index = n_iters - self.cur_index - 1;
            let base_iter = &mut self.base_iters[index];
            if let Some(key) = base_iter.next_back() {
                return Some((
                    key.clone(),
                    self.lru.peek(key.clone()).unwrap(),
                ));
            }

            if self.cur_index == n_iters - 1 {
                return None;
            }

            self.cur_index += 1;
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::utils::unique_id::UniqueID;

    use super::Lru;

    #[test]
    fn test_lru() {
        let mut lru = Lru::with_capacity(2, 20);

        for i in 0..10 {
            lru.insert(UniqueID(0, i), i);
        }

        for i in 0..10 {
            lru.insert(UniqueID(0, 10 + i), 10 + i);
        }

        assert_eq!(lru.len(), 10);

        for v in lru.iter() {
            println!("{}: {}", v.0, v.1);
        }

        for i in 0..10 {
            assert!(lru.peek(UniqueID(0, 10 + i)).is_some());
        }

        // 将10-19元素提升到L2队列
        for i in 0..10 {
            lru.access(UniqueID(0, 10 + i));
        }

        // 将20-29号元素插入到L1队列
        for i in 0..10 {
            lru.insert(UniqueID(0, 20 + i), 20 + i);
        }

        for v in lru.iter() {
            println!("{}: {}", v.0, v.1);
        }

        println!("===================");

        // 将20提升到L2, 将10淘汰到L1.
        // 将21提升到L2, 将21淘汰到L1.
        lru.access(UniqueID(0, 20));
        lru.access(UniqueID(0, 21));

        for v in lru.iter() {
            println!("{}: {}", v.0, v.1);
        }

        println!("===================");

        // 将30-39号元素插入到L1队列
        for i in 0..10 {
            lru.insert(UniqueID(0, 30 + i), 30 + i);
        }

        for v in lru.iter() {
            println!("{}: {}", v.0, v.1);
        }

        println!("===================");

        for v in lru.iter().rev() {
            println!("{}: {}", v.0, v.1);
        }

        println!("===================");

        // 删除10-19号元素 (如果存在).
        for i in 0..10 {
            lru.remove(UniqueID(0, 10 + i));
        }

        // 淘汰30-39号元素.
        for _ in 0..10 {
            lru.evict();
        }

        // 只剩下20和21号元素.
        assert_eq!(lru.len(), 2);
        assert!(lru.peek(UniqueID(0, 20)).is_some());
        assert!(lru.peek(UniqueID(0, 21)).is_some());
    }
}
